import torch
from server import Server
from client import Client
from utils import load_datasets, setup_logger

def main():
    # 实验配置
    config = {
        'num_clients': 100,
        'num_rounds': 100,
        'non_iid_degree': 0.1,
        'pruning_threshold': 0.05,
        'dbscan_eps': 0.3,
        'dbscan_min_samples': 5,
        'device': 'cuda' if torch.cuda.is_available() else 'cpu'
    }
    
    logger = setup_logger()
    logger.info("Starting MFL Framework Experiment")
    
    # 加载数据集
    train_datasets, test_datasets = load_datasets(config['num_clients'], config['non_iid_degree'])
    
    # 初始化服务器和客户端
    server = Server(config)
    clients = [Client(client_id, train_datasets[client_id], test_datasets[client_id], config) 
               for client_id in range(config['num_clients'])]
    
    # 联邦训练循环
    for round in range(config['num_rounds']):
        logger.info(f"Round {round+1}/{config['num_rounds']}")
        
        # 客户端本地训练
        for client in clients:
            client.local_train()
        
        # 服务器聚合与优化
        server.aggregate(clients)
        server.optimize_clusters()
        
        # 模型分发
        server.distribute_models(clients)
    
    # 最终评估
    server.evaluate(clients)
    logger.info("Experiment Completed")

if __name__ == "__main__":
    main()
